"""
C段，旁站，网站IP扫描
Class c network segment、side station、IP scan
"""

from common import *


class IpScan:
    def __init__(self, url, semaphore, side_station=False, save_path=None):
        """
        定义参数 Define parameters.
        :param url: 目标url    Target url
        :param semaphore: 并发数   Number of concurrent
        :param side_station: 是否扫描旁站     Whether scan side station
        :param save_path: 保存的文件路径   Save file path
        """

        self._url = url
        self._semaphore = semaphore
        self.side_station = side_station
        self._save_path = save_path
        # 扫描结果  Scanning result
        self.result = None
        # 输出表格  Form
        self.table = PrettyTable(['ip', 'state', 'port'])

    def _side_scan(self):
        """
        扫描旁站
        接口是webscan.cc的接口
        :return:
        """
        response_domain = requests.get("https://api.webscan.cc/?action=query&ip=" + self.scan_url)
        self.result = response_domain.json()
        success_color.print(f'{self._url}的旁站如下：')
        for i in self.result:
            success_color.print(f'[+]同IP域名：{i.get("domain")}')

    async def _net_scan(self, scan_ip):
        """
        扫描C段    Class c network segment scan
        :return:
        """

        async with self._semaphore:
            try:
                if await asyncio.wait_for(asyncio.open_connection(scan_ip, 80), timeout=0.3):
                    self.table.add_row([scan_ip, 'open', 80])
                elif await asyncio.wait_for(asyncio.open_connection(scan_ip, 443), timeout=0.3):
                    self.table.add_row([scan_ip, 'open', 443])
            except:
                pass

    def scan_options(self):
        """
        扫描选项    scan options
        :return:
        """
        # 正则匹配ip地址格式
        rez = re.compile(r'\d+\.\d+\.\d+\.')
        ip = rez.findall(self._url)
        if self.side_station:
            self._side_scan()
        else:
            pass
            if len(ip) > 0:
                scan_ip = ip[0]+"0"+"/24"
            else:
                # 访问接口获取IP地址
                req = requests.get("https://api.webscan.cc/?action=getip&domain=" + self._url, timeout=2)
                scan_ip = req.json().get('ip')
                # 正则匹配IP
                aa = rez.findall(scan_ip)
                scan_ip = aa[0]+"0"+"/24"
            ip_segment = ipaddress.IPv4Network(scan_ip).hosts()
            ip_list = []
            for i in ip_segment:
                i = str(i)
                ip_list.append(i)
            batch_run(func=self._net_scan, parma=ip_list,
                      semaphore=self._semaphore, table=self.table,
                      save_path=self._save_path, save_func=self.save_result)

    def save_result(self):
        """
        保存扫描结果  Save result data.
        :return:
        """
        if self._save_path:
            with open(self._save_path+'.csv', 'w', newline='') as fp:
                # 将表格转换为json字符串用json.loads读取出来
                # Load a table into a json string with json.loads.
                table = json.loads(self.table.get_json_string())
                # Take out the list head.
                # 取出表头
                header = table[0]
                table.remove(table[0])
                # 写入数据
                # Write data.
                writer = csv.DictWriter(fp, header)
                writer.writeheader()
                writer.writerows(table)
            success_color.print('保存成功\t', os.getcwd() + '\\' + self._save_path + '.csv')
